package cn.buhler.source;

import cn.buhler.jdbc.JDBCConfig;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.flume.*;
import org.apache.flume.conf.Configurable;
import org.apache.flume.conf.ConfigurationException;
import org.apache.flume.event.EventBuilder;
import org.apache.flume.source.AbstractSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @author shallwe_wang(xiaowei.wang@buhlergroup.com)
 *  2023年7月30日 下午1:30:39
 * 简单的自定义数据源
 */
public class JDBCSource extends AbstractSource implements Configurable, PollableSource {
    private static final Logger LOG = LoggerFactory.getLogger(JDBCSource.class);
    //轮询间隔
    private Integer pollInterval;
    //sql语句
    private String query;
    private String columns;

    /**
     * 是否选择所有的列发送到通道中
     */
    private boolean selectAllColumns;

    private HikariDataSource dataSource;

    /**
     * 线程正在运行的标记，用于通知线程资源的释放
     */
    private volatile boolean running;

    @Override
    public void configure(Context context) {
        try {
            JDBCConfig jdbcConfig = new JDBCConfig(context);
            this.dataSource = jdbcConfig.getDataSource();
        } catch (ConfigurationException e) {
            throw e;
        }
        query = context.getString("query");
        if(StrUtil.isEmpty(query)){
            throw new ConfigurationException(".query参数为空，用于同步的sql语句必须指定。");
        }
        columns = context.getString("columns", "*");
        if("*".equals(columns)) {
            this.selectAllColumns = true;
        }
        pollInterval = context.getInteger("poll.interval", 60);
    }


    @Override
    public synchronized void start() {
        running = true;
        super.start();
    }

    /**
     * 执行query参数提供的sql语句，把查询结果按照要求发送到Flume通道中
     * start后一直被Flume调用，需要
     * @return Status
     * @throws EventDeliveryException
     */
    @Override
    public Status process() throws EventDeliveryException {
        if (running) {
            List<Event> eventBatch = new ArrayList<>();
            try {
                List<Entity> results = Db.use(dataSource).query(query);

                for (Entity entity : results) {
                    //默认发送所有的列
                    Entity toSend = entity;
                    //如果不是要发送所有列
                    if (!selectAllColumns) {
                        toSend = new Entity();
                        String[] targetColumns = columns.split(",");
                        for (String column : targetColumns) {
                            if (entity.containsKey(column)) {
                                toSend.set(column, entity.get(column));
                            }
                        }
                    }
                    eventBatch.add(EventBuilder.withBody(toSend.toString().getBytes()));
                }
            } catch (SQLException e) {
                LOG.error("执行sql:{}读取数据发生错误.{}", query,e);
            }
            //批量发送数据到Flume通道中
            if (!eventBatch.isEmpty()) {
                getChannelProcessor().processEventBatch(eventBatch);
            }
        }
        /**
         * 线程休息指定的间隔
         */
        try {
            TimeUnit.SECONDS.sleep(pollInterval);
        } catch (InterruptedException e) {
            LOG.warn("读取数据的线程发生中断", e);
            Thread.currentThread().interrupt();

        }
        return Status.READY;
    }


    @Override
    public void stop() {
        running = false;
        try {
            if (this.dataSource != null) {
                this.dataSource.close();
                LOG.info("成功关闭连接池");
            }
        } catch (Exception e) {
            throw new FlumeException("关闭连接池失败", e);
        }
        super.stop();
    }

    /**
     * 暂时没有用到
     * @return
     */
    @Override
    public long getBackOffSleepIncrement() {
        return 0;
    }

    /**
     * 暂时没有用到
     * @return
     */
    @Override
    public long getMaxBackOffSleepInterval() {
        return 0;
    }
}
